# -*- coding: utf-8 -*-
# 这行指定了文件的编码格式为UTF-8，确保中文字符能够正确显示

# 导入必要的库
import json  # 用于处理JSON格式数据，如读取法律条文文件
import time  # 用于计时，比如记录索引加载时间
from pathlib import Path  # 提供面向对象的文件系统路径操作
from typing import List, Dict  # 类型提示，帮助IDE提供更好的代码补全

# 向量数据库和RAG相关库
import chromadb  # 向量数据库，用于存储和检索文本的向量表示
from llama_index.core import VectorStoreIndex, StorageContext, Settings  # RAG核心组件
from llama_index.core.schema import TextNode  # 文档的基本单位，每个法律条文一个节点
from llama_index.llms.huggingface import HuggingFaceLLM  # 对接HuggingFace大语言模型
from llama_index.embeddings.huggingface import HuggingFaceEmbedding  # 文本转向量
from llama_index.vector_stores.chroma import ChromaVectorStore  # llama_index与chromadb连接器
from llama_index.core import PromptTemplate  # 定义提示模板

# ================== 提示模板定义 ==================
# 定义问答模板，这是给大模型的指令和上下文格式
# {context_str}和{query_str}是占位符，运行时会被实际内容替换
QA_TEMPLATE = (
    "<|im_start|>system\n"
    "你是一个专业的法律助手，请严格根据以下法律条文回答问题：\n"
    "相关法律条文：\n{context_str}\n<|im_end|>\n"
    "<|im_start|>user\n{query_str}<|im_end|>\n"
    "<|im_start|>assistant\n"
)

# 创建提示模板对象，后续会用于查询引擎
response_template = PromptTemplate(QA_TEMPLATE)


# ================== 配置类 ==================
class Config:
    """
    配置类：集中管理所有配置参数，方便修改和维护
    类似于项目的"设置中心"
    """
    # 模型路径配置
    EMBED_MODEL_PATH = r"D:/models/BAAI_bge-small-zh-v1.5"  # 嵌入模型路径，用于文本转向量
    LLM_MODEL_PATH = r"D:\models\models\deepseek-ai\DeepSeek-R1-Distill-Qwen-1___5B"  # 大语言模型路径

    # 数据路径配置
    DATA_DIR = "./law_data"  # 原始法律条文JSON文件存放目录
    VECTOR_DB_DIR = "./law_chroma_db"  # 向量数据库持久化目录
    PERSIST_DIR = "./law_storage"  # 索引持久化目录

    # 向量数据库配置
    COLLECTION_NAME = "chinese_labor_laws"  # 集合名称，标识这组法律数据
    TOP_K = 3  # 检索时返回的最相关条目数量


# ================== 模型初始化 ==================
def init_models():
    """
    初始化嵌入模型和大语言模型
    返回: (embed_model, llm) - 嵌入模型和语言模型实例
    """
    # 初始化嵌入模型：将文本转换为向量表示
    # 例如："劳动合同" -> [0.1, 0.2, 0.3, ...] (512维向量)
    embed_model = HuggingFaceEmbedding(
        model_name=Config.EMBED_MODEL_PATH,  # 模型路径
        device='cuda' if hasattr(Settings, 'device') else 'cpu'  # 优先使用GPU加速
    )

    # 初始化大语言模型：用于生成回答
    # 例如：根据检索到的条文，生成对"试用期最长多久"的回答
    llm = HuggingFaceLLM(
        model_name=Config.LLM_MODEL_PATH,  # 模型路径
        tokenizer_name=Config.LLM_MODEL_PATH,  # 分词器路径
        device_map="auto",  # 自动选择设备(GPU/CPU)
        tokenizer_kwargs={"trust_remote_code": True},  # 信任模型的远程代码
        generate_kwargs={"temperature": 0.3}  # 温度参数：控制生成随机性，越低越确定性
    )

    # 将模型设置为全局默认，这样其他组件会自动使用这些模型
    Settings.embed_model = embed_model
    Settings.llm = llm

    # 验证嵌入模型是否正常工作
    test_embedding = embed_model.get_text_embedding("测试文本")
    print(f"✅ 嵌入模型验证成功，向量维度：{len(test_embedding)}")

    return embed_model, llm


# ================== 数据处理模块 ==================
def load_and_validate_json_files(data_dir: str) -> List[Dict]:
    """
    加载并验证JSON格式的法律文件
    支持两种格式：列表格式和字典格式

    参数:
        data_dir: 数据目录路径

    返回:
        List[Dict]: 标准化后的数据列表
    """
    # 查找目录中所有JSON文件
    json_files = list(Path(data_dir).glob("*.json"))
    # 确保至少找到一个JSON文件
    assert json_files, f"❌ 错误：在目录 {data_dir} 中未找到任何JSON文件"

    all_data = []  # 存储所有标准化后的数据

    for json_file in json_files:
        print(f"📁 正在处理文件: {json_file.name}")
        with open(json_file, 'r', encoding='utf-8') as f:
            try:
                data = json.load(f)  # 解析JSON文件

                # ===== 格式1：列表格式 =====
                # 示例: [{"法律条文1": "内容1"}, {"法律条文2": "内容2"}]
                if isinstance(data, list):
                    print(f"  📋 检测到列表格式，包含 {len(data)} 个条目")
                    for item in data:
                        # 验证每个条目是否为字典
                        if not isinstance(item, dict):
                            raise ValueError(f"文件 {json_file.name} 包含非字典元素")
                        # 验证字典中的每个值是否为字符串
                        for k, v in item.items():
                            if not isinstance(v, str):
                                raise ValueError(f"文件 {json_file.name} 中键 '{k}' 的值不是字符串")
                    # 将列表格式转换为标准格式
                    all_data.extend({
                                        "content": item,  # 法律条文内容
                                        "metadata": {"source": json_file.name}  # 元数据：来源文件
                                    } for item in data)

                # ===== 格式2：字典格式 =====
                # 示例: {"法律条文1": "内容1", "法律条文2": "内容2"}
                elif isinstance(data, dict):
                    print(f"  📚 检测到字典格式，包含 {len(data)} 个键值对")
                    for full_title, content in data.items():
                        # 验证内容是否为字符串
                        if not isinstance(content, str):
                            raise ValueError(f"文件 {json_file.name} 中键 '{full_title}' 的值不是字符串")
                        # 将字典格式转换为标准格式
                        all_data.append({
                            "content": {full_title: content},  # 包装成字典格式
                            "metadata": {"source": json_file.name}
                        })
                else:
                    raise ValueError(f"文件 {json_file.name} 根元素应为列表或字典")

            except Exception as e:
                raise RuntimeError(f"❌ 加载文件 {json_file} 失败: {str(e)}")

    print(f"✅ 成功加载 {len(all_data)} 个法律文件条目")
    return all_data


def create_nodes(raw_data: List[Dict]) -> List[TextNode]:
    """
    将原始数据转换为TextNode对象
    每个法律条文转换为一个TextNode，包含内容和元数据

    参数:
        raw_data: 标准化后的原始数据

    返回:
        List[TextNode]: TextNode对象列表
    """
    nodes = []

    for entry in raw_data:
        law_dict = entry["content"]  # 法律条文字典
        source_file = entry["metadata"]["source"]  # 来源文件名

        # 遍历字典中的每个法律条文
        for full_title, content in law_dict.items():
            # 生成稳定ID：确保每次运行生成相同的ID
            # 格式: "文件名::完整标题"
            node_id = f"{source_file}::{full_title}"

            # 分割标题，提取法律名称和条款号
            # 示例: "中华人民共和国劳动法 第一条" ->
            #       法律名称="中华人民共和国劳动法", 条款号="第一条"
            parts = full_title.split(" ", 1)  # 按第一个空格分割
            law_name = parts[0] if len(parts) > 0 else "未知法律"
            article = parts[1] if len(parts) > 1 else "未知条款"

            # 创建TextNode对象
            node = TextNode(
                text=content,  # 条文具体内容
                id_=node_id,  # 稳定ID
                metadata={
                    "law_name": law_name,  # 法律名称
                    "article": article,  # 条款号
                    "full_title": full_title,  # 完整标题
                    "source_file": source_file,  # 来源文件
                    "content_type": "legal_article"  # 内容类型标识
                }
            )
            nodes.append(node)

    print(f"✅ 生成 {len(nodes)} 个文本节点")
    if nodes:
        print(f"  示例节点ID：{nodes[0].id_}")
    return nodes


# ================== 向量存储模块 ==================
def init_vector_store(nodes: List[TextNode]) -> VectorStoreIndex:
    """
    初始化向量存储，创建或加载索引
    这是RAG系统的核心组件

    参数:
        nodes: TextNode对象列表

    返回:
        VectorStoreIndex: 向量索引对象
    """
    # 创建持久化的ChromaDB客户端
    # 数据将保存在指定目录，下次启动可直接加载
    chroma_client = chromadb.PersistentClient(path=Config.VECTOR_DB_DIR)

    # 获取或创建集合（类似于数据库中的表）
    chroma_collection = chroma_client.get_or_create_collection(
        name=Config.COLLECTION_NAME,
        metadata={
            # 向量距离度量配置 (hnsw:space)
            "hnsw:space": "cosine",
            # 可选值:
            #   "cosine" - 余弦相似度：最适合文本相似性搜索，衡量向量方向而非大小
            #   "l2"     - 欧几里得距离：衡量向量间的直线距离，适合图像、坐标等数据
            #   "ip"     - 内积相似度：衡量向量的点积，在某些机器学习模型中常用
            # HNSW 索引构建参数
            "hnsw:construction_ef": 200,
            # 构建时的 ef 参数 (默认: 200)
            # - 控制图构建过程中的搜索范围，值越大构建质量越高但速度越慢
            # - 影响索引的准确性和构建时间，建议范围: 100-500

            "hnsw:search_ef": 100,
            # 搜索时的 ef 参数 (默认: 20)
            # - 控制搜索时的候选集大小，值越大搜索结果越准确但搜索速度越慢
            # - 在生产环境中需要根据准确性和延迟要求进行权衡

            "hnsw:M": 16,
            # 每个节点的最大连接数 (默认: 16)
            # - 控制图的连通性，值越大图越密集，搜索精度越高但内存占用越大
            # - 通常设置在 8-32 之间，文本数据常用 16

            # 性能优化参数
            "hnsw:num_threads": 4,
            # 并行线程数 (默认: CPU核心数)
            # - 控制构建和搜索时使用的线程数量
            # - 在多核CPU上可以显著提升性能

            "hnsw:batch_size": 100,
            # 批量处理大小 (默认: 100)
            # - 批量插入或搜索时的批次大小
            # - 较大的批次可以提高吞吐量但增加内存使用

            "hnsw:sync_threshold": 1000,
            # 同步阈值 (默认: 1000)
            # - 内存中累积多少更改后同步到磁盘
            # - 值越大性能越好，但故障时可能丢失更多数据

            # 自定义元数据 (可选)
            "description": "文档语义搜索集合",  # 集合描述，便于管理
            "version": "1.0",  # 集合版本
            "created_by": "AI团队",  # 创建者信息
            "data_type": "text_embeddings",  # 数据类型说明
        })

    # 创建存储上下文，连接到ChromaDB
    storage_context = StorageContext.from_defaults(
        vector_store=ChromaVectorStore(chroma_collection=chroma_collection)
    )

    # 判断是否需要新建索引
    if chroma_collection.count() == 0 and nodes is not None:
        print(f"🆕 创建新索引（{len(nodes)}个节点）...")

        # 将节点添加到文档存储中
        storage_context.docstore.add_documents(nodes)

        # 创建向量索引：这是最耗时的步骤
        # 会对每个节点的文本进行向量化，并建立索引结构
        index = VectorStoreIndex(
            nodes,
            storage_context=storage_context,
            show_progress=True  # 显示进度条
        )

        # 双重持久化保障，确保数据不会丢失
        storage_context.persist(persist_dir=Config.PERSIST_DIR)
        index.storage_context.persist(persist_dir=Config.PERSIST_DIR)

    else:
        # 如果集合不为空，则加载已有索引
        print("📂 加载已有索引...")
        storage_context = StorageContext.from_defaults(
            persist_dir=Config.PERSIST_DIR,
            vector_store=ChromaVectorStore(chroma_collection=chroma_collection)
        )
        index = VectorStoreIndex.from_vector_store(
            storage_context.vector_store,
            storage_context=storage_context,
            embed_model=Settings.embed_model
        )

    # 存储验证：确保数据正确加载
    print("\n🔍 存储验证结果：")
    doc_count = len(storage_context.docstore.docs)
    print(f"  📊 DocStore记录数：{doc_count}")

    if doc_count > 0:
        sample_key = next(iter(storage_context.docstore.docs.keys()))
        print(f"  🔖 示例节点ID：{sample_key}")
    else:
        print("  ⚠️  警告：文档存储为空！")

    return index


# ================== 主程序 ==================
def main():
    """
    主函数：程序入口，协调各个模块的工作流程
    """
    print("🚀 开始初始化法律智能助手...")

    # 1. 初始化模型
    print("\n" + "=" * 50)
    print("步骤1: 初始化模型")
    print("=" * 50)
    embed_model, llm = init_models()

    # 2. 数据处理：仅当需要更新数据时执行
    print("\n" + "=" * 50)
    print("步骤2: 数据处理")
    print("=" * 50)
    if not Path(Config.VECTOR_DB_DIR).exists():
        print("📦 初始化数据（首次运行或更新数据）...")
        raw_data = load_and_validate_json_files(Config.DATA_DIR)  # 加载原始数据
        nodes = create_nodes(raw_data)  # 转换为TextNode
    else:
        print("📂 使用已有数据...")
        nodes = None  # 已有数据，不需要重新处理

    # 3. 初始化向量存储
    print("\n" + "=" * 50)
    print("步骤3: 向量存储初始化")
    print("=" * 50)
    start_time = time.time()  # 开始计时
    index = init_vector_store(nodes)
    load_time = time.time() - start_time
    print(f"✅ 索引加载完成，耗时：{load_time:.2f}秒")

    # 4. 创建查询引擎
    print("\n" + "=" * 50)
    print("步骤4: 创建查询引擎")
    print("=" * 50)
    query_engine = index.as_query_engine(
        similarity_top_k=Config.TOP_K,  # 返回最相关的3个结果
        text_qa_template=response_template,  # 使用自定义提示模板
        verbose=True  # 显示详细日志
    )
    print("✅ 查询引擎准备就绪")

    # 5. 交互式问答循环
    print("\n" + "=" * 50)
    print("步骤5: 进入交互问答")
    print("=" * 50)
    print("💡 您可以开始提问了！输入 'q' 退出程序")

    while True:
        try:
            # 获取用户输入
            question = input("\n❓ 请输入劳动法相关问题: ").strip()

            # 退出条件
            if question.lower() == 'q':
                print("👋 感谢使用，再见！")
                break
            if not question:
                print("⚠️  问题不能为空，请重新输入")
                continue

            # 执行查询
            print("🔍 正在检索相关法律条文...")
            response = query_engine.query(question)

            # 显示回答结果
            print(f"\n🤖 智能助手回答：")
            print(f"   {response.response}")

            # 显示支持依据（检索到的相关法律条文）
            print(f"\n📚 支持依据（检索到 {len(response.source_nodes)} 个相关条文）:")
            for idx, node in enumerate(response.source_nodes, 1):
                meta = node.metadata
                print(f"\n  [{idx}] {meta['full_title']}")
                print(f"     来源文件：{meta['source_file']}")
                print(f"     法律名称：{meta['law_name']}")
                print(f"     条款内容：{node.text[:80]}...")  # 显示前80个字符
                print(f"     相关度得分：{node.score:.4f}")  # 相似度分数

        except KeyboardInterrupt:
            print("\n👋 用户中断，程序退出")
            break
        except Exception as e:
            print(f"❌ 处理问题时出错: {e}")


# 程序入口点
if __name__ == "__main__":
    """
    当直接运行此脚本时执行main函数
    如果被其他模块导入，则不执行
    """
    main()
